import os
import traceback
from copy import deepcopy

from process_v18 import messages_to_function_call
# from util.query_openai_api import LLMClient

from util.query_qwen_vllm import LLMClient

if __name__ == '__main__':
    import dotenv
    import json
    from concurrent.futures import ThreadPoolExecutor, as_completed

    dotenv.load_dotenv()
    llm = LLMClient()

    with open("data/合成数据决赛赛题.jsonl", "r", encoding="utf-8") as f:
        # 载入jsonl
        test_date = [json.loads(line) for line in f]

    print("测试集载入完成: ", len(test_date))

    results_data = deepcopy(test_date)

    def process_data(i):
        print("正在生成第", i, "个数据")
        # if i>100:
        #     return [[] for _ in range(len(test_date[i]["user_messages"]))]

        targets = [[] for _ in range(len(test_date))]
        for r in range(6):  # Retry up to 6 times
            try:
                targets = messages_to_function_call(llm, test_date[i]["user_messages"], test_date[i]['apis'], retry=False if r==0 else True)
            except Exception as e:
                print("<UNK>", i, e)
                continue

            # 如果至少有一次函数调用，那么就算成功。
            if any(len(target) > 0 for target in targets):
                break

        results_data[i]['targets']=targets
        print(targets)
        return targets

    with ThreadPoolExecutor(max_workers=30) as executor:
        future_to_index = {executor.submit(process_data, i): i for i in range(len(test_date))}
        for future in as_completed(future_to_index, timeout=600000):  # 设置超时为120秒
            index = future_to_index[future]
            try:
                result = future.result()
                # dataset.extend(result)
            except Exception as e:
                print(f"处理第 {index} 个数据时发生错误: {e}")
                traceback.print_exc()

    print("完成: ")

    for i in range(len(results_data)):
        try:
            assert len(results_data[i]['targets']) == len(test_date[i]['user_messages'])
            assert results_data[i]['id'] == test_date[i]['id']
        except Exception as e:
            print(f"<UNK> {i} <UNK>",e)

    with open("result/final_qwen_0320.jsonl", "w", encoding="utf-8") as f:
        for l in results_data:
            # ，每个json对象一行
            s = json.dumps(l, ensure_ascii=False)
            f.writelines(s + "\n")
